import cv2
import time
from road_class import RoadClass
from IPython.display import display, Image, clear_output
from pydub import AudioSegment
import simpleaudio as sa

audio_map = {
    "level": AudioSegment.from_wav("/home/Code/RoboticDog/resource/pre_audio_file/level.wav"),
    "stairs": AudioSegment.from_wav("/home/Code/RoboticDog/resource/pre_audio_file/stairs.wav"),
    "overleap": AudioSegment.from_wav("/home/Code/RoboticDog/resource/pre_audio_file/overleap_audio.wav"),
    "slope": AudioSegment.from_wav("/home/Code/RoboticDog/resource/pre_audio_file/slope_audio.wav"),
}

def show_camera_feed(camera_index=0, window_width=640, window_height=480):
    """
    在Jupyter Notebook中实时显示摄像头视频
    
    参数:
        camera_index: 摄像头索引(默认为0)
        window_width: 显示窗口宽度
        window_height: 显示窗口高度
    """
    box_x, box_y = 420, 0
    box_width, box_height = 1080, 1080

    rc = RoadClass()
    status = "None"

    cap = cv2.VideoCapture(camera_index)

    
    if not cap.isOpened():
        print("无法打开摄像头")
        return

    try:
        while True:
            # 清除之前的输出
            clear_output(wait=True)
            
            ret, frame = cap.read()
            if not ret:
                print("无法获取帧")
                break
                
            # 将帧从BGR转换为RGB
            img = frame
            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            
            cv2.rectangle(frame, (box_x, box_y), 
                         (box_x + box_width, box_y + box_height), 
                         (0, 255, 0), 2)

            _, encoded = cv2.imencode('.jpg', frame)
    
            # 显示新帧
            display(Image(data=encoded.tobytes()))

            results = rc.infer(img)

            if status != results:
                status = results
                if status != 'None':
                    print("1111111111111111111111111111", status)
                    play_obj = sa.play_buffer(audio_map[status].raw_data, num_channels=audio_map[status].channels, bytes_per_sample=audio_map[status].sample_width, sample_rate=audio_map[status].frame_rate)
            
            # 添加小延迟，避免过度占用资源
            time.sleep(0.001)

    except KeyboardInterrupt:
        print("摄像头已停止")
    finally:
        cap.release()

# 使用示例
show_camera_feed()